import os
import shutil

# 原始数据集文件夹路径
source_dir = r'E:\mnist_jpg'

# 目标数据集文件夹路径
target_dir = r'E:\mnist'

# 创建目标文件夹
os.makedirs(os.path.join(target_dir, 'training'), exist_ok=True)
os.makedirs(os.path.join(target_dir, 'test'), exist_ok=True)

# 获取源目录中的所有文件
files = os.listdir(source_dir)

# 遍历文件并根据文件名分配到train和test文件夹
for file in files:
    # 检查文件是否符合train或test的格式
    if file.startswith('train_'):
        # 移动到train文件夹
        shutil.move(os.path.join(source_dir, file), os.path.join(target_dir, 'train', file))
    elif file.startswith('test_'):
        # 移动到test文件夹
        shutil.move(os.path.join(source_dir, file), os.path.join(target_dir, 'test', file))

# 输出分类完成信息
print("文件已成功分类！")
